"""
使用克里金重采样偏差，校正降尺度结果
思路：
计算低分辨率下原始粗分辨率SM每个像元对应高分辨率像元平均值的偏差
得到粗分辨率下的偏差二维数组，使用克里金重采样到高分辨率
然后将偏差加到高分辨率的降尺度SM中得到偏差移除的结果
"""
import os.path
from pykrige.ok import OrdinaryKriging
from scipy.interpolate import griddata
from osgeo import gdal
import numpy as np


# 写入tif文件
def write_tif(file_path, data, geotrans, projection, nodata, gdal_type):
    driver = gdal.GetDriverByName("GTiff")
    rows, cols = data.shape
    dataset = driver.Create(file_path, cols, rows, 1, gdal_type)
    dataset.SetGeoTransform(geotrans)
    # 定义投影
    # prj = osr.SpatialReference()
    # prj.ImportFromEPSG(4326)
    # dataset.SetProjection(prj.ExportToWkt())
    dataset.SetProjection(projection)
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(nodata)
    dataset.FlushCache()
    print("写入成功：{}".format(file_path))


# 获取低分辨率像元对应所有高分辨率像元值
def get_high_res_pixels(high_res_geo_transform, high_res_data, low_res_geo_transform, x_low_res_pixel, y_low_res_pixel, low_res_to_high_res_ratio):
    # 计算低分辨率影像像元在高分辨率影像中的范围
    x_low_res_geo, y_low_res_geo = pixel2coord(low_res_geo_transform, x_low_res_pixel, y_low_res_pixel)
    x_low_res_geo = x_low_res_geo + low_res_geo_transform[1] / 2  # 左右偏移
    y_low_res_geo = y_low_res_geo + low_res_geo_transform[5] / 2  # 上下偏移

    # 计算行列号范围
    high_res_x, high_res_y = coord2pixel(high_res_geo_transform, x_low_res_geo, y_low_res_geo)

    high_res_x_min = int(high_res_x - low_res_to_high_res_ratio / 2)
    high_res_x_max = int(high_res_x + low_res_to_high_res_ratio / 2)
    high_res_y_min = int(high_res_y - low_res_to_high_res_ratio / 2)
    high_res_y_max = int(high_res_y + low_res_to_high_res_ratio / 2)

    # 排除边缘像元的影响
    if high_res_x_min < 0 and high_res_x_max > 0:
        high_res_x_min = 0
    if high_res_x_min < high_res_data.shape[1] < high_res_x_max:
        high_res_x_max = high_res_data.shape[1]
    if high_res_y_min < 0 and high_res_y_max > 0:
        high_res_y_min = 0
    if high_res_y_min < high_res_data.shape[0] < high_res_y_max:
        high_res_y_max = high_res_data.shape[0]

    # 获取范围内的像元值
    if 0 <= high_res_x_min <= high_res_data.shape[1] and 0 <= high_res_y_min <= high_res_data.shape[0]:
        high_res_pixels = high_res_data[high_res_y_min:high_res_y_max, high_res_x_min:high_res_x_max]
        return high_res_pixels.astype(np.float32)
    else:
        return []


# 根据像元行列号计算经纬度
def pixel2coord(geo_transform, x, y):
    x_origin, pixel_width, x_rotation, y_origin, y_rotation, pixel_height = geo_transform
    x_geo = x_origin + pixel_width * x
    y_geo = y_origin + pixel_height * y
    return x_geo, y_geo


# 根据经纬度计算像元行列号
def coord2pixel(geo_transform, x_geo, y_geo):
    x = int((x_geo - geo_transform[0]) / geo_transform[1])
    y = int((y_geo - geo_transform[3]) / geo_transform[5])
    return x, y

# 移除偏差
def remove_bias(orign_path, downscal_path, out_path):
    orign_dataset = gdal.Open(orign_path)
    downscal_dataset = gdal.Open(downscal_path)
    orign_band = orign_dataset.GetRasterBand(1)
    downscal_band = downscal_dataset.GetRasterBand(1)
    orign_data = orign_band.ReadAsArray()
    downscal_data = downscal_band.ReadAsArray()
    orign_data = orign_data.astype(np.float32)
    downscal_data = downscal_data.astype(np.float32)
    orign_geo = orign_dataset.GetGeoTransform()
    downscal_geo = downscal_dataset.GetGeoTransform()
    downscal_projection = downscal_dataset.GetProjection()
    orign_data[orign_data == -9999] = np.nan
    downscal_data[downscal_data == -9999] = np.nan

    low_bias = np.zeros(orign_data.shape, dtype=np.float32)
    high_bias = np.zeros(downscal_data.shape, dtype=np.float32)

    # 像元比例
    res_ratio = int(orign_geo[1] / downscal_geo[1])
    for i in range(orign_data.shape[0]):
        for j in range(orign_data.shape[1]):
            high_pixel = get_high_res_pixels(downscal_geo, downscal_data, orign_geo, j, i, res_ratio)
            # print(high_pixel)
            mean = np.nanmean(high_pixel)
            low_bias[i][j] = orign_data[i][j] - mean
    # 进行克里金插值
    coords = []
    values = []
    for i in range(orign_data.shape[0]):
        for j in range(orign_data.shape[1]):
            if not np.isnan(low_bias[i][j]):
                coords.append((i * res_ratio + int(res_ratio / 2), j * res_ratio + int(res_ratio / 2)))
                values.append(low_bias[i][j])
    # 定义新的高分辨率网格
    high_x = np.linspace(0, downscal_data.shape[0] - 1, downscal_data.shape[0])
    high_y = np.linspace(0, downscal_data.shape[1] - 1, downscal_data.shape[1])
    high_coords = np.meshgrid(high_x, high_y)
    high_coords = np.vstack([high_coords[0].flatten(), high_coords[1].flatten()]).T

    high_values = griddata(coords, values, high_coords, method='cubic')
    for i in range(high_values.shape[0]):
        high_bias[int(high_coords[i][0]), int(high_coords[i][1])] = high_values[i]
    high_bias[np.isnan(high_bias)] = 0
    print(high_bias)
    rm_downscal = downscal_data + high_bias
    # print(rm_downscal)
    output_path = out_path + "/" + downscal_path.split("\\")[-1].replace(".tif", "_rmbias.tif")
    write_tif(output_path, rm_downscal, downscal_geo, downscal_projection, -9999, gdal.GDT_Float32)

if __name__ == '__main__':
    orign = r"G:\test\downscal_result\orign"
    downscal = r"G:\test\downscal_result\downscal"
    out = r"G:\test\downscal_result\downscal_rmbias"
    if not os.path.exists(out):
        os.makedirs(out)
    orign_list = [os.path.join(orign, i) for i in os.listdir(orign)]
    downscal_list = [os.path.join(downscal, i) for i in os.listdir(downscal)]
    for orign_path, downscal_path in zip(orign_list, downscal_list):
        remove_bias(orign_path, downscal_path, out)